import { BlockEnum } from '@/components/workflow/types'
import { produce } from 'immer'
import { useMemo } from 'react'
import { NODES_EXTRA_DATA } from '../constants'

export const useNodesExtraData = () => {
  return produce(NODES_EXTRA_DATA, draft => {
    Object.keys(draft).forEach(key => {
      draft[key as BlockEnum].availablePrevNodes = draft[key as BlockEnum].getAvailablePrevNodes()
      draft[key as BlockEnum].availableNextNodes = draft[key as BlockEnum].getAvailableNextNodes()
    })
  })
}

export const useAvailableBlocks = (nodeType?: BlockEnum, isInIteration?: boolean) => {
  const nodesExtraData = useNodesExtraData()
  const availablePrevBlocks = useMemo(() => {
    if (!nodeType) return []
    return nodesExtraData[nodeType].availablePrevNodes || []
  }, [nodeType, nodesExtraData])

  const availableNextBlocks = useMemo(() => {
    if (!nodeType) return []
    return nodesExtraData[nodeType].availableNextNodes || []
  }, [nodeType, nodesExtraData])

  return useMemo(() => {
    return {
      availablePrevBlocks: availablePrevBlocks.filter(_nType => {
        if (isInIteration && nodeType !== BlockEnum.End 
          &&  nodeType !== BlockEnum.LoopEnd  
          && nodeType !==  BlockEnum.IterationEnd && nodeType !==  BlockEnum.TextConcatenation)   return false
        return true
      }),
      // 不采用dify设计逻辑，直接通过限制入口来控制连接，不限制出口
      availableNextBlocks: availableNextBlocks.filter(nType => {
        // if (isInIteration && (nType === BlockEnum.Iteration || nType === BlockEnum.End)) return false
        return true
      })
    }
  }, [isInIteration, availablePrevBlocks, availableNextBlocks])
}
